from pyalgotrade import strategy
from pyalgotrade.barfeed import yahoofeed
from pyalgotrade.technical import ma
from pyalgotrade.stratanalyzer import returns
from pyalgotrade import plotter


class MyStrategy(strategy.BacktestingStrategy):
    def __init__(self, feed, instrument, smaPeriod):
        super(MyStrategy, self).__init__(feed, 1000)
        self.__position = None
        self.__instrument = instrument
        # We'll use adjusted close values instead of regular close values.
        self.setUseAdjustedValues(True)
        self.__sma = ma.SMA(feed[instrument].getPriceDataSeries(), smaPeriod)
        self.count = 0

    def onEnterOk(self, position):
        execInfo = position.getEntryOrder().getExecutionInfo()
        self.info("BUY at $%.2f" % (execInfo.getPrice()))

    def onEnterCanceled(self, position):
        self.__position = None

    def onExitOk(self, position):
        execInfo = position.getExitOrder().getExecutionInfo()
        self.info("SELL at $%.2f" % (execInfo.getPrice()))
        self.__position = None

    def onExitCanceled(self, position):
        # If the exit was canceled, re-submit it.
        self.__position.exitMarket()

    def onBars(self, bars):
        self.count += 1
        bar = bars[self.__instrument]
        dt = bar.getDateTime()
        close = bar.getPrice()
        print(str(self.count) + " " + str(dt) + " " + str(close))

        if self.count == 1:
            self.__position = self.enterLong(self.__instrument, 10, True)


        # Wait for enough bars to be available to calculate a SMA.
        #if self.__sma[-1] is None:
        #    return

        bar = bars[self.__instrument]
        # If a position was not opened, check if we should enter a long position.
        #if self.__position is None:
        #    if bar.getPrice() > self.__sma[-1]:
                # Enter a buy market order for 10 shares. The order is good till canceled.
        #        self.__position = self.enterLong(self.__instrument, 10, True)
        # Check if we have to exit the position.
        #elif bar.getPrice() < self.__sma[-1] and not self.__position.exitActive():
        #    self.__position.exitMarket()


def run_strategy(smaPeriod):
    # Load the yahoo feed from the CSV file
    path = "D:/Data/PyAlgoTradeData/Yahoo/"
    feed = yahoofeed.Feed()
    feed.addBarsFromCSV("000001.SZ", path + "000001.SZ.csv")

    # Evaluate the strategy with the feed.
    myStrategy = MyStrategy(feed, "000001.SZ", smaPeriod)

    # Attach a returns analyzers to the strategy.
    returnsAnalyzer = returns.Returns()
    myStrategy.attachAnalyzer(returnsAnalyzer)

    # Attach the plotter to the strategy.
    plt = plotter.StrategyPlotter(myStrategy)
    # Include the SMA in the instrument's subplot to get it displayed along with the closing prices.
    # plt.getInstrumentSubplot("000001.SZ").addDataSeries("SMA", myStrategy.getSMA())
    # Plot the simple returns on each bar.
    #plt.getOrCreateSubplot("returns").addDataSeries("Simple returns", returnsAnalyzer.getReturns())

    myStrategy.run()
    print("Final portfolio value: $%.2f" % myStrategy.getBroker().getEquity())

    # Plot the strategy.
    plt.plot()

run_strategy(15)